package cn.edu.ncu.housinghub.config;

import cn.edu.ncu.housinghub.entity.user.User;
import cn.edu.ncu.housinghub.utils.common.TokenUtil;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

/**
 * WebSocket握手拦截器，用于身份验证
 */
@Component
public class WebSocketAuthInterceptor implements HandshakeInterceptor {

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
            
            // 从请求头获取token
            String token = servletRequest.getServletRequest().getHeader("info");
            
            // 如果没有token，尝试从查询参数获取
            if (token == null || token.isEmpty()) {
                token = servletRequest.getServletRequest().getParameter("token");
            }
            
            if (token != null && !token.isEmpty()) {
                try {
                    // 验证token并获取用户信息
                    User user = TokenUtil.getUserByToken(token);
                    if (user != null) {
                        // 将用户信息存储到WebSocket会话属性中
                        attributes.put("userId", user.getId());
                        attributes.put("username", user.getUsername());
                        attributes.put("token", token);
                        System.out.println("WebSocket握手成功，用户ID: " + user.getId());
                        return true;
                    }
                } catch (Exception e) {
                    System.err.println("WebSocket token验证失败: " + e.getMessage());
                }
            }
            
            System.err.println("WebSocket握手失败：无效的token");
        }
        
        return false; // 拒绝连接
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        if (exception != null) {
            System.err.println("WebSocket握手后发生异常: " + exception.getMessage());
        }
    }
}
